/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains stellarmodel, a class for reading a stellar model FITS file as
/// generated by convertstellarmodel. \ingroup sfrhist

#include "model.h"
#include "preferences.h"
#include <memory> 
#include <fstream> 
#include "CCfits/CCfits"
#include "blitz-fits.h"
#include "fits-utilities.h"
#include "misc.h"
#include "sphgrid.h"
#include "boost/lambda/lambda.hpp"
#include "interpolatort.h"
#include "constants.h"
#include "vecops.h"
#include "boost/foreach.hpp"
#include "boost/graph/graph_traits.hpp"
#include "boost/graph/adjacency_list.hpp"
#include <boost/graph/connected_components.hpp>

using namespace std;
using namespace CCfits;
using namespace blitz;



/** Reads the stellar SEDs from a FITS file. */
mcrx::simple_stellarmodel::simple_stellarmodel (Preferences& p) :
  prefs(p), n_gt_max_age(0), n_lt_min_age(0), n_gt_max_met(0), n_lt_min_met(0)
{
  // The first thing we should do is to open the FITS file containing
  // the stellar model specified in keyword "stellarmodelfile"
  auto_ptr<FITS> input_file;
  if (prefs.defined("stellarmodelfile")) {
    file_name = 
      word_expand(prefs.getValue("stellarmodelfile", string ()))[0];
    try {
      input_file.reset(new FITS (file_name, Read, "SED", true));
    }
    catch (...) {
      cerr << "Error: Can't open stellar model file: "
	   << file_name << endl;
      exit (1);
    }
  }
  else {
    cerr << "Error: No keyword \"stellarmodelfile\" defined" <<endl;
    exit (1);
  }
  cout << "Reading stellar model file: " << file_name <<endl;

  // Check that the file is actually a stellar model file
  string file_type;
  input_file->pHDU().readKey ("FILETYPE", file_type);
  if (file_type != "STELLARMODEL") {
    cerr << "Error: File " << file_type << " is not a stellar model file"
	 <<endl;
    exit (1);
  }

  ExtHDU & parameters = open_HDU(*input_file,"STELLARMODEL");
  parameters.readKey("logflux",  log_flux);
  // reference_mass has to be a float because readKey can't read doubles
  parameters.readKey("refmass",  reference_mass_);
  parameters.readKey("powerunit",  units_["luminosity"] );
  parameters.readKey("timeunit",  units_["time"] );
  parameters.readKey("lengthunit",  units_["wavelength"]);
  parameters.readKey("massunit",  units_["mass"]);
  
  // Check that the units are what we expect
  if ((units().get("luminosity") != "W") || (units().get("time") != "yr") ||
      (units().get("wavelength") != "m") || (units().get("mass") != "Msun")) {
    cout << "Warning: Nonstandard units in stellar model file\n";
    cout << units().get("luminosity") << ' ';
    cout << units().get("time") << ' ';
    cout << units().get("wavelength") << ' ';
    cout << units().get("mass") << '\n';
  }

  // Read the SED
  ExtHDU & spectrum = open_HDU(*input_file, "SED");

  // If we have a column called "time" then it's the old-style file
  // for only one metallicity where the SED is in a table.  If not,
  // it's the new-style file where the SED is in a 3-D data cube.
  try {
    Column & sed_c =
      spectrum.column((log_flux?  string ("log "): string ("")) + "L_lambda");
    array_2 temp;
    read (sed_c, temp);
    sed.resize(temp.rows(), 1, temp.columns() );
    sed (Range::all (), 0, Range::all ()) = temp;
    Column & time_c = spectrum.column("time" );
    assert (time_c.rows() == n_times ());
    time_c.read (times, 1, n_times ());

    T_float z;
    parameters.readKey("track_metallicity", z);
    metallicities.push_back(z );

    ExtHDU& l = open_HDU (*input_file, "LAMBDA");
    l.column("lambda").read( lambdas, 1, n_lambda () );

  }
  catch (ExtHDU::WrongExtensionType&) {
    // read SED from data cube
    read (spectrum, sed);
    try {
      // try to read current mass fraction
      ExtHDU & mass_hdu = open_HDU(*input_file, "current_mass");
      read (mass_hdu, current_mass);
    }
    catch (FITS::NoSuchHDU&) {
      // Not present
      cout << "Mass loss data not present in stellar model file"<<endl;
    }

    // read times and metallicities
    ExtHDU& tm = open_HDU (*input_file, "AXES");
    tm.column("time" ).read(times, 1, n_times () );
    tm.column("metallicity").read(metallicities, 1, n_metallicities () );
    tm.column("lambda").read( lambdas, 1, n_lambda () );
  }
  
  // deal with log flux and reference mass
  if (log_flux)
    sed = pow (10., sed)/reference_mass_ ;
  else  
    sed/= reference_mass_;
  // update these to reflect the current values
  log_flux = false;
  reference_mass_ = 1;

  // make sure the SED extends over the requested range, if not pad
  // with zero on the short end and a 1/lambda^4 decline on the long
  // end.

  // We choose the points extend ACROSS the extended range in order to
  // avoid problems with resampling.
  bool extend_front=false;
  bool extend_back=false;
  int s; // first wavelength of existing sed to keep
  int e; // last wavelength of existing sed to keep
  vector<T_float> new_lambda;

  if(p.getValue("min_wavelength", T_float()) < lambdas.front()) {
    // extension asked for, add an extra value to front of wavelength vector
    new_lambda.push_back(p.getValue("min_wavelength", T_float())*0.99);
    extend_front=true;
    s = 0;
  }
  else
    // truncation asked for, find min entry
    s = lower_bound (lambdas.begin(), lambdas.end(),
		     prefs.getValue ("min_wavelength", T_float ())) -
      lambdas.begin() - 1; // -1 to extend across range

  if(p.getValue("max_wavelength", T_float()) > lambdas.back()) {
    // extension asked for, add an extra value to back of wavelength vector
    copy(lambdas.begin()+s, lambdas.end(), back_inserter(new_lambda));
    new_lambda.push_back(p.getValue("max_wavelength", T_float()));
    extend_back=true;
    e = sed.extent(thirdDim)-1;
  }
  else {
    // truncation asked for, find max entry (no -1 here so we extend
    // to max_wavelength)
    e = lower_bound (lambdas.begin(), lambdas.end(),
		     prefs.getValue ("max_wavelength", T_float ())) -
      lambdas.begin();
    copy(lambdas.begin()+s, lambdas.begin()+e+1, back_inserter(new_lambda));
  }

  array_3 new_sed(n_times(), n_metallicities(), new_lambda.size());
  new_sed=blitz::quiet_NaN(T_float());
  // copy existing SED into new array
  new_sed(Range::all(), Range::all(), 
	  Range(extend_front?1:0, e-s + (extend_front?1:0) ) ) = 
    sed(Range::all(), Range::all(), Range(s,e));

  if(extend_front)
    new_sed(Range::all(), Range::all(), 0) = blitz::tiny(T_float());
  if(extend_back)
    // calculate end point as L_lambda ~ 1/lambda^4, assuming we are on
    // the Rayleigh-Jeans tail of the SED
    new_sed(Range::all(), Range::all(), new_lambda.size()-1) =
      new_sed(Range::all(), Range::all(), new_lambda.size()-2)*
      pow(new_lambda[new_lambda.size()-2]/new_lambda.back(), 4);

  swap(lambdas, new_lambda);
  assert (lambdas.size() == new_sed.extent(thirdDim));
  sed.reference(new_sed);

  ASSERT_ALL(sed>=0);
  ASSERT_ALL(sed<blitz::huge(T_float()));

  middle_times = calculate_middle(times);
  middle_metallicities = calculate_middle(metallicities);

  alambda_.reference(reference_vector(lambdas));

  calculate_L_bol();
}

void
mcrx::simple_stellarmodel::calculate_L_bol()
{
  // calculate bolometric luminosity by integrating sed.
  L_bol.resize(n_times (), n_metallicities () );
  for(int i=0;i<n_times();++i)
    for(int j=0;j<n_metallicities();++j) {
      L_bol(i,j) = integrate_quantity(sed(i,j,Range::all()),
				      lambda());
      assert(L_bol(i,j)<blitz::huge(T_float()));
    }
}


void
mcrx::simple_stellarmodel::
write_fits_parameters (const string& output_file_name) const
{
  FITS output_file(output_file_name, Write);
  auto_ptr<FITS> input_file (new FITS (file_name, Read));

  ExtHDU& st=open_HDU(*input_file,"STELLARMODEL");
  output_file.copy(st) ;

  // write LAMBDA HDU
  ExtHDU*lam =output_file.addTable("LAMBDA", 0);
  lam->addColumn(Tdouble, "lambda", 1, units().get("wavelength"));
  lam->column("lambda" ).write(lambdas, 1);
}


mcrx::array_1 mcrx::simple_stellarmodel::get_SED(star_particle_set& s, 
						 int index,
						 T_float time) const
{
  // get particle age and metallicity
  const T_float age = time - s.tf(index);
  const T_float metallicity = s.z(index);

  // Return the actual SED, correctly accounting for particle mass
  return array_1(get_SED(age, metallicity) * s.mcr(index));
}


mcrx::array_1 mcrx::simple_stellarmodel::get_SED(T_float age,
						 T_float metallicity) const
{
  if(age<0) {
      cerr << "WARNING: Asking for SED with negative age: " << age << endl;
    }
     
  // find the time index. This is the bin before the first bin for
  // which the bin midpoint is not less than the age.
  const int t =
    lower_bound(middle_times.begin(), middle_times.end(), age) - 
    middle_times.begin() - 1;
  assert (t >= 0);
  assert (t < sed.rows());

  if(age<times.front()) {
    DEBUG(1,cerr << "WARNING: Asking for SED with age < minimum age: " << age
      << " < " << times.front() << endl;);
    n_lt_min_age++;
  }
  if(age>times.back()){
    DEBUG(1,cerr << "WARNING: Asking for SED with age > maximum age: " << age
      << " > " << times.back() << endl;);
    n_gt_max_age++;
  }

  // find the metallicity index
  const int m = 
    lower_bound(middle_metallicities.begin(), middle_metallicities.end(),
		metallicity) - middle_metallicities.begin() - 1;
  assert (m >= 0);
  assert (m < sed.columns());

  if(metallicity<metallicities.front()){
    DEBUG(1,cerr << "WARNING: Asking for SED with metallicity < minimum "
      << "metallicity: " << metallicity << " < " << metallicities.front()
      << endl;);
    n_lt_min_met++;
  }

  if(metallicity>metallicities.back()){
    DEBUG(1,cerr << "WARNING: Asking for SED with metallicity > maximum "
      "metallicity: " << metallicity << " > " << metallicities.back() << endl;);
    n_gt_max_met++;
  }

  // Note that this returns the mass-normalized SED, since there is no
  // mass info in this call.
  return sed (t, m, Range::all ());
}


void
mcrx::simple_stellarmodel::
calculate_SEDs(const std::vector<star_particle_ptr>& particles, 
	       T_float time) const
{
  cout << "Calculating SEDs for SB99 particles" << endl;

  // This is simple, just loop over all particles and call get_SED.
  BOOST_FOREACH(const star_particle_ptr& p, particles) {
    const int n= p.num();
    p.sed()(n, Range::all()) = get_SED(p.set(), n, time);
  }

  if(n_gt_max_age + n_lt_min_age > 0){
    cout << "WARNING: " << n_gt_max_age+n_lt_min_age
         << " particles outside  the age range of the SB99 templates ("
         << n_gt_max_age << " > max, " << n_lt_min_age << " < min)."
         << " If this number is large it is a problem." << endl;
  }
  if(n_gt_max_met + n_lt_min_met > 0){
  cout << "WARNING: " << n_gt_max_met+n_lt_min_met
       << " particles outside the metallicity range of the SB99 templates ("
       << n_gt_max_met << " > max, " << n_lt_min_met << " < min)."
       << " If this number is large it is a problem." << endl;
  }
}


double mcrx::simple_stellarmodel::get_mass_fraction (T_float age,
						     T_float metallicity) const
{
  // If model doesn't include mass loss data, we assume no mass loss
  if(current_mass.size()==0) return 1;

  // find the time index
  const int t =
    find_if (middle_times.begin(), middle_times.end(),
	     bind2nd (greater<double> (), age)) -
    middle_times.begin() - 1;
  assert (t >= 0);
  assert (t < sed.rows());

  // find the metallicity index
  const int m = 
    find_if (middle_metallicities.begin(), middle_metallicities.end(),
	     bind2nd (greater<double> (), metallicity)) -
    middle_metallicities.begin() - 1;
  assert (m >= 0);
  assert (m < sed.columns());

  const double cm=current_mass (t, m);
  assert(cm>=0);
  assert(cm<=1);

  return current_mass (t, m);
}


void mcrx::simple_stellarmodel::resample(const array_1& new_lambda)
{
  // check if we are upsampling or downsampling. If upsampling, we use
  // mcrx::resample with a logarithmic resampling, because that works
  // better. For downsampling, we use mcrx::subsample, which preserves
  // the integral.
  const bool upsample = 
    (log10(lambda()(lambda().size()-1))-log10(lambda()(0)))/(lambda().size()-1) >
    ((log10(new_lambda(new_lambda.size()-1))-log10(new_lambda(0)))/
     (new_lambda.size()-1));
  if (upsample) 
    cout << "Upsampling SEDs onto wavelength grid\n";
  else
    cout << "Downsampling SEDs onto wavelength grid\n";

  array_3 new_sed(sed.extent(firstDim), sed.extent(secondDim),
		  new_lambda.size());

  if (upsample)
    for(int i=0; i <sed.extent(firstDim); ++i)
      for(int j=0; j <sed.extent(secondDim); ++j) {
	new_sed(i,j,Range::all()) = mcrx::resample(sed(i,j,Range::all()),
						   lambda(),
						   new_lambda, false, true);
      }
  else
    for(int i=0; i <sed.extent(firstDim); ++i)
      for(int j=0; j <sed.extent(secondDim); ++j) {
	new_sed(i,j,Range::all()) = mcrx::subsample(sed(i,j,Range::all()),
						    lambda(),
						    new_lambda);
    }

  // swap in the new SED array
  sed.reference(new_sed);
  lambdas.assign(new_lambda.begin(), new_lambda.end());

  alambda_.reference(reference_vector(lambdas));

  // L_bol *should* not have changed much but better be self-consistent
  calculate_L_bol();
}


mcrx::mappings_stellarmodel::mappings_stellarmodel(Preferences& p) :
  prefs(p), axes(nvar_), middle_axes(nvar_),
  fixed_mcl_(prefs.defined("mappings_mcl")), set_particle_radius_(false),
  n_gt_max_age(0), n_lt_min_age(0), n_gt_max_met(0), n_lt_min_met(0),
  n_gt_max_S(0), n_lt_min_S(0), n_gt_max_P(0), n_lt_min_P(0)
{
  if(fixed_mcl_) {
    Mcl_ = log10(prefs.getValue("mappings_mcl", T_float()));
    assert(!prefs.defined("mappings_S"));
  }
  else {
    S_ = log10(prefs.getValue("mappings_S", T_float()));
    assert(!prefs.defined("mappings_mcl"));
  }

  pdr_f_ = prefs.getValue("mappings_pdr_fraction", T_float ());

  use_teff_ = p.defined("mappings_use_teff") && 
    p.getValue("mappings_use_teff", bool());

  do_clustering_ = p.defined("cluster_mappings_particles") &&
    p.getValue("cluster_mappings_particles", bool());

  stupid();

  if(p.defined("mappings_pdr_file")) {
    const string pdrfile(word_expand(p.getValue("mappings_pdr_file", string()))[0]);
    cout << "Reading MAPPINGS PDR radii from file " << pdrfile << endl;
    ifstream rfile(pdrfile.c_str());
    assert(rfile.good());
    r_pdr_.resize(shape(6,5));
    for (int S=0; S<6;++S)
      for (int Pk=0; Pk<5;++Pk)
	rfile >> r_pdr_(S,Pk);
    r_pdr_*=units::convert("cm","kpc");
    cout << r_pdr_;
    set_particle_radius_ = p.defined("mappings_set_radius") &&
      p.getValue("mappings_set_radius", bool());
  }
  else {
    if(do_clustering_) {
      cerr << "WARNING: Must define mappings_pdr_file in order to use "
           << "cluster_mappings_particles. Turning clustering off."
	   << endl;
      do_clustering_ = false;
    }
  }
}


void mcrx::mappings_stellarmodel::stupid()
{
  using namespace boost::lambda;
  const string input = 
    word_expand(prefs.getValue("mappings_sed_file", string()))[0];
  cout << "Reading MAPPINGS SED from file " << input << endl;

  auto_ptr<FITS> input_file;
  try {
    input_file.reset(new FITS(input, Read )); 
  }
  catch (FITS::CantOpen) {
    cerr << "Failed opening MAPPINGS file: " << input << endl;
    throw;
  }

  // look for file version. If version doesn't exist is the old file
  // that has the incorrect COVER axis.
  string version;
  try {
    open_HDU(*input_file, 2).readKey ("VERSION",version);
  }
  catch (CCfits::HDU::NoSuchKeyword&) {
  }  

  T_float refmass;
  try {
    open_HDU(*input_file, 2).readKey ("REFMASS", refmass);
  }
  catch (CCfits::HDU::NoSuchKeyword&) {
    cout << "Warning: MAPPINGS reference mass not found in file, assuming 1e7 Msun\n";
    refmass=1e7;
  }

  Array< T_float, nvar_> temp;
  read(open_HDU(*input_file, 2), temp);
  ExtHDU& axes_HDU=open_HDU (*input_file, 1);

  // to get in increasing wavelength order
  // and transpose (Z, S, P/k, 0/1, l) to get axes to be (S, P/k, Z,
  // l, 0/1). We put the HII/PDR axes last so the numbering of the
  // others don't change from temp to the sed array.
  temp.transposeSelf(secondDim, thirdDim, firstDim, fifthDim, fourthDim);
  temp.reverseSelf(fourthDim);


  // make sure that HII region is index 1 and PDR is 0. (Note that the
  // COVER values in Brent's old files from 2007 are incorrect. While
  // the axes values are 0,1, the HII region is actually in index 0.)
  // Thus, if there is any version, it's later and the axes are
  // correct.
  if(version!="") {
    cout << "MAPPINGS file version " << version << endl;
    array_1 covering;
    read(axes_HDU.column("COVER"), covering);
    if(covering(0)==0){
      // HII region is in index 0, reverse axis
      assert(covering(1)>0);
      temp.reverseSelf(fifthDim);
    }
  }
  else
    cout << "Unversioned MAPPINGS file" << endl;

  // AXES vector is (S, P/k, Z, t, l)

  // MAKE UP time axis for now, since the mappings sed is
  // time-averaged. mappings sed extends from 0 to 10Myr one element
  // at 1e7 should mean it applies up to that time
  axes[t_axis].resize(2);
  axes[t_axis][0]=0;
  axes[t_axis][1]=1e7;

  axes_HDU.column("SPARAM").read (axes [s_axis], 1, temp.extent(0));
  axes_HDU.column("PRESSURE").read (axes [p_axis], 1, temp.extent(1 ));
  axes_HDU.column("METAL").read (axes [z_axis], 1, temp.extent(2));
  // convert metallicity from units of solar to mass fraction
  transform (axes [z_axis].begin(), axes [z_axis].end(), 
	     axes [z_axis].begin(), _1 *= 0.02);

  axes_HDU.column("WAVE").read (axes [l_axis], 1, temp.extent(3 ));
  // flip so axis is in order of increasing wavelength
  reverse(axes[l_axis].begin(), axes [l_axis].end()); 

  // convert to m
  transform (axes [l_axis].begin(), axes [l_axis].end(), axes [l_axis].begin(),
  	     _1 *= 1e-6);

  cout << "  Using MAPPINGS PDR fraction " << pdr_f_ << endl;
  // Copy data into SED array so we have the indices correct, while
  // applying clearing timescale.  furthermore, brents files are L_nu,
  // not L_lambda.  We need to convert by doing L_lambda = L_nu *
  // c/lambda^2 and convert from erg/s to W and take care of 1e7 Msun
  // mass norm of the mappings SEDs.
  const array_1 lambda(reference_vector(axes[l_axis]));

  if (fixed_mcl_)
    cout << "  Using MAPPINGS cluster mass " << Mcl_ << endl;
  else
    cout << "  Using MAPPINGS S-parameter " << S_ << endl;

  sed.resize(axes[s_axis].size(), axes[p_axis].size(), 
	     axes[z_axis].size(), axes[t_axis].size(), axes[l_axis].size());

  sed (Range::all(), Range::all(), Range::all(), 0, Range::all()) =
    (pdr_f_*temp(Range::all(), Range::all(), Range::all(), 
		Range::all(), 0) +
     (1-pdr_f_)*temp(Range::all(), Range::all(), Range::all(), 
		    Range::all(), 1)) *
    (constants::c0*units::convert("erg","J")/
     (refmass*lambda(tensor::l)*lambda(tensor::l)));

  sed (Range::all(), Range::all(), Range::all(), 1, Range::all()) =
    sed (Range::all(), Range::all(), Range::all(), 0, Range::all());

  cout << "Reading MAPPINGS SED with size " << temp.size() << endl;
  cout << "axes sizes " << axes [0].size() << ' '
       << axes [1].size() << ' '
       << axes [2].size() << ' '
       << axes [3].size() << ' '
       << axes [4].size() << '\n' << temp.shape() << endl;

  // kill wavelengths outside of the sb99 range
  const int s = lower_bound (axes [l_axis].begin(), axes [l_axis].end(),
			     prefs.getValue ("min_wavelength", T_float ())) -
    axes [l_axis].begin(); 
  // -1 so that we don't exceed the requested wavelength
  const int e = lower_bound (axes [l_axis].begin(), axes [l_axis].end(),
			     prefs.getValue ("max_wavelength", T_float ())) -
    axes [l_axis].begin()-1;

  cout << "  Using wavelength range " << s << " - " << e << endl;
  assert(e>=0);
  sed.reference(sed (Range::all (), Range::all (), Range::all (),
		     Range::all (), Range (s, e)));
  axes [l_axis].erase(axes [l_axis].begin()+e+1, axes [l_axis].end() ); 
  axes [l_axis].erase(axes [l_axis].begin(), axes [l_axis].begin() +s);
  assert (axes [l_axis].size() == sed.extent(l_axis));

  middle_axes.resize(axes.size());
  for(int i = 0; i < axes.size(); ++i)
    middle_axes[i] = calculate_middle(axes[i]);

  alambda_.reference(reference_vector(axes[l_axis]));
};


mcrx::array_2 mcrx::mappings_stellarmodel::read_brent_file(const string& fn)
{
  ifstream f(fn.c_str());
  assert(f.good());
  string s;
  getline(f,s);
  getline(f,s);
  getline(f,s);
  getline(f,s);
  getline(f,s);
  char c;
  f.get(c);
  vector<T_float> temp((istream_iterator<T_float>(f)), 
		       istream_iterator<T_float>());
  assert(f.eof());
  assert((temp.size()>0)&&(temp.size()%2==0));
  array_2 arr(&temp[0], shape(temp.size()/2, 2), duplicateData);

  return arr;
}

/** Reads the .source file with all times in it that brent supplied. */
mcrx::array_2 mcrx::mappings_stellarmodel::read_brent_file2(const string& fn)
{
  ifstream f(fn.c_str());
  assert(f.good());
  string s;
  getline(f,s);
  getline(f,s);
  getline(f,s);
  getline(f,s);
  vector<T_float> temp((istream_iterator<T_float>(f)), 
		       istream_iterator<T_float>());
  assert(f.eof());
  assert((temp.size()>0)&&(temp.size()%24==0));
  array_2 arr(&temp[0], shape(temp.size()/24, 24), duplicateData);
  // brent has ordered in increasing energy, so need to flip
  arr.reverseSelf(firstDim);
  arr(Range::all(), 2) *= 1e-6;
  // furthermore, brents files are L_nu, not L_lambda.
  // we need to convert by doing L_lambda = L_nu * c/lambda^2
  // and convert from erg/s to W
  array_2 templ(arr(Range::all(), Range(3,arr.extent(secondDim)-1)));
  array_1 lambda(arr(Range::all(), 2));
  templ = templ(tensor::i, tensor::j)*constants::c0*1e-7*1e-5/(lambda(tensor::i)*lambda(tensor::i));

  return arr;
}


mcrx::T_float
mcrx::mappings_stellarmodel::check_pdr_mass(T_float rho_gas, T_float z_gas,
					    T_float Pk, T_float age, 
					    T_float radius, T_float S, 
					    T_float mass) const
{
  if (r_pdr_.size()==0) {
    cout << "PDR radius data not available, can't compute statistics" << endl;
    return 0;
  }

  const T_float r_pdr = get_pdr_radius(Pk,S);

  const T_float ncl = mass/pow(10.,Mcl_);
  const T_float m_pdr = 1e9*r_pdr*r_pdr*pdr_f_;
  const T_float rho_pdr = 2.4e8/r_pdr;
  const T_float r_s = pow( 3*m_pdr / (4*rho_gas*constants::pi), 1./3);

  cout << "MAPPINGS mass double-counting check:" <<
    "\n\tparticle mass:\t" << mass << 
    "\n\tMappings cluster mass:\t" << pow(10.,Mcl_) <<
    "\n\tNumber of clusters in particle:\t" << ncl << 
    "\n\tMappings PDR fraction:\t" << pdr_f_ <<
    "\n\tPDR mass:\t" << m_pdr << 
    "\n\tPDR density:\t" << rho_pdr << 
    "\n\tgas density in cell:\t" << rho_gas << 
    "\n\tgas metallicity in cell:\t" << z_gas << 
    "\n\tparticle radius:\t" << radius << 
    "\n\tPDR radius:\t" << r_pdr <<
    "\n\tradius needed to get PDR mass:\t" 
       << r_s <<
    "\n\tApprox double-counted optical depth:\t" 
       << 5.86e-6*rho_gas*(r_s-radius)*z_gas*0.4 <<
    "\n\tHII region intensity log10(U)=\t" << -1.62+S <<
    "\n\tPDR intensity log10(U)=\t" << -2.56+2*S/3+log10(Pk)/3
       << endl;

  return r_s;
}




mcrx::array_1 mcrx::mappings_stellarmodel::get_SED(star_particle_set& s, 
						   int index,
						   T_float time) const
{
  // get particle age and metallicity
  const T_float age = time - s.tf(index);
  const T_float metallicity = s.z(index);

  const pair<T_float,pair<T_float, T_float> > hydro = pressure_functor(s.pos(index));
  const T_float Pk = hydro.first;
  const T_float rho_gas = hydro.second.first;
  const T_float z_gas = hydro.second.second;

  // if we have fixed mcl we calculate S, but if we have fixed S we
  // just use it.
  const T_float S = fixed_mcl_? (3*Mcl_ + 2*log10(Pk))/5 : S_;

  // print pdr mass check stuff
  const T_float r_s = check_pdr_mass(rho_gas, z_gas, Pk, age, s.h(index), S, s.mcr(index));

  // set the particle radius to r_s if that option is set.
  if(set_particle_radius_ && rho_gas>blitz::tiny(T_float()))
    s.h(index)=r_s;

  // We are now returning the actual SED, not the mass-normalized one,
  // so we need to multiply by the (creation) mass of the particle
  // here.
  return array_1(get_SED(age, metallicity, Pk, S) * s.mcr(index));
}
 
mcrx::T_float mcrx::mappings_stellarmodel::get_pdr_radius(T_float Pk,
							  T_float S) const
{
  using namespace boost::lambda;

  int Sindex = find_if (middle_axes[s_axis].begin(), middle_axes[s_axis].end(),
			_1 > S) -
    middle_axes[s_axis].begin() - 1;
  int Pindex = find_if (middle_axes[p_axis].begin(), middle_axes[p_axis].end(),
			_1 > log10(Pk)) -
    middle_axes[p_axis].begin() - 1;

  return r_pdr_(Sindex,Pindex);
}

mcrx::array_1 mcrx::mappings_stellarmodel::get_SED(T_float age,
						   T_float metallicity,
						   T_float Pk,
						   T_float Sparam) const
{
  using namespace boost::lambda;

  TinyVector<double, nvar_> x;
  x[s_axis] = Sparam;
  x [p_axis] = log10(Pk);
  x [z_axis] = metallicity;
  x [t_axis] = age;
  
  cout << "Asking for SED with S, pressure, Z, age: " << x << endl;
  if(age<0) {
      cerr << "WARNING: Asking for SED with negative age: " << age << endl;
    }

  // Check if age is outside range of Groves models
  if(age > axes[t_axis].back()){
      DEBUG(1,cerr << "WARNING: Asking for SED with age > maximum age in Groves models (" << axes[t_axis].back() << ")" << endl;);
      n_gt_max_age++;
  }
  if(age < axes[t_axis].front()){
      DEBUG(1,cerr << "WARNING: Asking for SED with age < minimum age in Groves models (" << axes[t_axis].front() << ")" << endl;);
      n_lt_min_age++;
  }

  // Check if compactness parameter is outside range of Groves models
  if(Sparam > axes[s_axis].back()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with S > maximum S in Groves models (" << axes[s_axis].back() << ")" << endl;);
      n_gt_max_S++;
  }
  if(Sparam < axes[s_axis].front()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with S < minimum S in Groves models (" << axes[s_axis].front() << ")" << endl;);
      n_lt_min_S++;
  }

  // Check if pressure is outside range of Groves models
  if(log10(Pk) > axes[p_axis].back()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with pressure > maximum pressure in Groves models (" << axes[p_axis].back() << ")" << endl;);
      n_gt_max_P++;
  }
  if(log10(Pk) < axes[p_axis].front()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with pressure < minimum pressure in Groves models (" << axes[p_axis].front() << ")" << endl;);
      n_lt_min_P++;
  }

  // Check if metallicity is outside range of Groves models
  if(metallicity > axes[z_axis].back()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with Z > maximum Z in Groves models (" << axes[z_axis].back() << ")" << endl;);
      n_gt_max_met++;
  }
  if(metallicity < axes[z_axis].front()) {
      DEBUG(1,cerr << "WARNING: Asking for SED with Z < minimum Z in Groves models (" << axes[z_axis].front() << ")" << endl;);
      n_lt_min_met++;
  }

  TinyVector<int, nvar_> sed_index;
  // find the indices
  for (int i=0; i < nvar_; ++i) {
    if (i == l_axis)
      continue;
    
    sed_index[i] = find_if (middle_axes[i].begin(), middle_axes[i].end(),
			    _1 > x[i]) -
      middle_axes[i].begin() - 1;
    assert(sed_index[i] >= 0);
    assert(sed_index[i] <  middle_axes[i].size());
  }

  cout << " Returning SED element: " << sed_index << endl;
  // Returns mass-normalized SED since we have no mass info here
  return sed (sed_index[0], sed_index[1], sed_index[2],
	      sed_index[3], Range::all ());
}

void
mcrx::mappings_stellarmodel::
calculate_SEDs(const std::vector<star_particle_ptr>& particles, 
	       T_float time) const
{
  if(!do_clustering_) {
    // we don't want to cluster the mappings particles. Just use the
    // simple method
    BOOST_FOREACH(const star_particle_ptr& p, particles) {
      const int n= p.num();
      // get_SED already accounts for mass
      p.sed()(n, Range::all()) = get_SED(p.set(), n, time);
    }
    if(n_gt_max_age+n_lt_min_age > 0){
      cout << "WARNING: " << n_gt_max_age+n_lt_min_age
           << " particles outside  the age range of the Groves templates ("
           << n_gt_max_age << " > max, " << n_lt_min_age
           << " < min). If this number is large it is a problem." << endl;
    }
    if(n_gt_max_met+n_lt_min_met > 0){
      cout << "WARNING: " << n_gt_max_met+n_lt_min_met
           << " particles outside the metallicity range of the Groves templates ("
           << n_gt_max_met << " > max, " << n_lt_min_met
           << " < min). If this number is large it is a problem." << endl;
    }
    if(n_gt_max_S+n_lt_min_S > 0){
      cout << "WARNING: " << n_gt_max_S+n_lt_min_S
           << " particles outside the compactness parameter range of the Groves "
           << "templates (" << n_gt_max_S << " > max, " << n_lt_min_S
           << " < min). If this number is large it is a problem." << endl;
    }
    if(n_gt_max_P+n_lt_min_P > 0){
      cout << "WARNING: " << n_gt_max_P+n_lt_min_P
           << " particles outside the pressure range of the Groves templates ("
           << n_gt_max_P << " > max, " << n_lt_min_P
           << " < min). If this number is large it is a problem." << endl;
    }

    return;
  }

  // move on to the clustering algorithm

  using namespace boost;
  typedef adjacency_list<vecS,vecS,undirectedS> Graph;
  typedef graph_traits<Graph>::vertex_descriptor Vertex;
  typedef graph_traits<Graph>::edge_descriptor Edge;
  typedef property_map<Graph, vertex_index_t>::type IndexMap;

  const T_float link_factor = 1;
  const int n=particles.size();

  // Calculate PDR radii for particles assuming a cluster mass of
  // their actual mass
  cout << "Calculating PDR radii" << endl;
  vector<T_float> r_pdr;
  vector<T_float> Pk;
  for(int i=0; i<n; ++i) {
    const pair<T_float,pair<T_float, T_float> > hydro = 
      pressure_functor(particles[i].set().pos(particles[i].num()));
    const T_float ppk = hydro.first;
    const T_float rho_gas = hydro.second.first;
    const T_float z_gas = hydro.second.second;

    const T_float mass = particles[i].set().m(particles[i].num());
    // calculate S-param assuming cluster mass=particle mass
    const T_float S = (3*log10(mass) + 2*log10(ppk))/5;
    // and get pdr radius
    const T_float r = get_pdr_radius(ppk, S);
    r_pdr.push_back(r);
    Pk.push_back(ppk);

    cout << "Particle " << i << ": Pk=" << ppk<< ", S=" << S << ", r_pdr= " << r << endl;
  }

  cout << "Adding " << n << " MAPPINGS particle to graph" << endl;

  // we create the graph by collecting all the edges in a vector. The
  // numbers of the vertices are the entries inte the particles
  // vector. For now we create an edge in a dumb n^2 way if the
  // particles overlap.
  vector<pair<int, int> > edges;
  for(int v1=0; v1<n-1; ++v1)
    for(int v2=v1+1; v2<n; ++v2) {
      const star_particle_ptr& p1 = particles[v1];
      const star_particle_ptr& p2 = particles[v2];
      const vec3d pos1= p1.set().pos(p1.num());
      const vec3d pos2= p2.set().pos(p2.num());
      const T_float r2= dot(pos1-pos2, pos1-pos2);
      const T_float pr= r_pdr[v1]+r_pdr[v2];
      if (r2< pr*pr*link_factor*link_factor) {
	edges.push_back(make_pair(v1,v2));
      }
    }

  Graph g(edges.begin(), edges.end(), n);
  cout << edges.size() << " edges in graph." << endl;
  cout << "Searching for clusters by computing connected components in graph" << endl;
  vector<int> component(n);
  const int num = connected_components(g, &component[0]);

  cout << num << " clusters of MAPPINGS particles:" << endl;

  // Now order the particles by cluster
  vector<vector<int> > clusters(num);
  for(int i=0; i<n; ++i)
    clusters[component[i]].push_back(i);

  ofstream junk("clusters.txt");
  for(int i=0; i<num; ++i) {
    cout << "Cluster " << i << " contains " << clusters[i].size() << " : ";

    T_float m_cl=0;
    if(clusters[i].size()>1) {
      // add up mass in particles in the cluster
      for(int j=0; j<clusters[i].size(); ++j) {
	const star_particle_ptr& p= particles[clusters[i][j]];
	m_cl += p.set().m(p.num());
	const vec3d pos=p.set().pos(p.num());
	junk << pos[0] << '\t' << pos[1] << '\t' << pos[2] << '\t' << p.set().h(p.num()) << '\t' << r_pdr[clusters[i][j]] << '\n';
	cout << clusters[i][j] << " ";
      }
      cout << '\n';
      cout << "\tCluster mass is " << m_cl << endl;
    }
    else {
      // cluster is a single particle. then we use the set mcl.
      m_cl = pow(10,Mcl_);
      cout << "\tUsing fixed cluster mass of " << m_cl << endl;
      const star_particle_ptr& p= particles[clusters[i][0]];
      const vec3d pos=p.set().pos(p.num());
      junk << pos[0] << '\t' << pos[1] << '\t' << pos[2] << '\t' << p.set().h(p.num()) << '\t' << r_pdr[clusters[i][0]] << '\n';
    }

    junk << "\n\n";

    // pass2: calculate SED now that we know the cluster mass
    for(int j=0; j<clusters[i].size(); ++j) {
      const star_particle_ptr& p= particles[clusters[i][j]];
      const T_float S = (3*log10(m_cl) + 2*log10(Pk[clusters[i][j]]))/5;
      const T_float age = time - p.set().tf(p.num());

      p.sed()(p.num(), Range::all()) = get_SED(age, 
					       p.set().z(p.num()),
					       Pk[clusters[i][j]],
					       S) * p.set().mcr(p.num());
    }
  }

  if(n_gt_max_age+n_lt_min_age > 0){
    cout << "WARNING: " << n_gt_max_age+n_lt_min_age
         << " particles outside  the age range of the Groves templates ("
         << n_gt_max_age << " > max, " << n_lt_min_age
         << " < min). If this number is large it is a problem." << endl;
  }
  if(n_gt_max_met+n_lt_min_met > 0){
    cout << "WARNING: " << n_gt_max_met+n_lt_min_met
         << " particles outside the metallicity range of the Groves templates ("
         << n_gt_max_met << " > max, " << n_lt_min_met
         << " < min). If this number is large it is a problem" << endl;
  }
  if(n_gt_max_S+n_lt_min_S > 0){
    cout << "WARNING: " << n_gt_max_S+n_lt_min_S
         << " particles outside the compactness parameter range of the Groves "
         << "templates (" << n_gt_max_S << " > max, " << n_lt_min_S
         << " < min). If this number is large it is a problem." << endl;
  }
  if(n_gt_max_P+n_lt_min_P > 0){
    cout << "WARNING: " << n_gt_max_P+n_lt_min_P
         << " particles outside the pressure range of the Groves templates ("
         << n_gt_max_P << " > max, " << n_lt_min_P
         << " < min). If this number is large it is a problem." << endl;
  }
}


mcrx::mappings_sb99_model::mappings_sb99_model(Preferences& p) :
  mappings(p), sb99(p) 
{
  // in general, the mappings model will have higher wavelength
  // resolution than the sb99 model, so we resample the sb99 model
  sb99.resample(mappings.lambda());

  // see if we actually want the mappings SEDs (this is different from
  // not specifying a mappings_sed_file, because here you still get
  // the SEDs sampled to the mappings wavelength grid.
  use_mappings = p.defined("use_mappings_seds")? p.getValue("use_mappings_seds", bool()) : true;

  check_consistency(mappings.units(), sb99.units(), true);
}


mcrx::array_1 mcrx::mappings_sb99_model::get_SED(star_particle_set& s, 
						 int index,
						 T_float time) const
{
  // see if stellar age is within mappings times, then use those
  const T_float age = time - s.tf(index);

  if ((age < mappings.oldest() ) && use_mappings)
    return mappings.get_SED(s, index, time);
  else
    return sb99.get_SED(s, index, time);
}

mcrx::array_1 mcrx::mappings_sb99_model::get_SED(T_float age, 
						 T_float metallicity, 
						 T_float Pk, 
						 T_float Sparam) const
{
  // see if stellar age is within mappings times, then use those
  if (age < mappings.oldest() )
    return mappings.get_SED(age, metallicity, Pk, Sparam);
  else
    return sb99.get_SED(age, metallicity);
}

void
mcrx::mappings_sb99_model::
calculate_SEDs(const vector<star_particle_ptr>& particles, 
	       T_float time) const
{
  // We partition the particles into those that should go to the sb99
  // model and to the mappings model, based on their time
  vector<star_particle_ptr> sb99parts, mapparts;

  BOOST_FOREACH(const star_particle_ptr& p, particles) {
    const T_float age = time - p.set().tf(p.num());

    if (use_mappings && (age < mappings.oldest())) 
      mapparts.push_back(p);
    else
      sb99parts.push_back(p);
  }
  mappings.calculate_SEDs(mapparts, time);
  sb99.calculate_SEDs(sb99parts, time);
}

